from langchain_chroma import Chroma
from langchain_community.document_loaders import TextLoader

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import CharacterTextSplitter

loader = TextLoader("D:\\dev\\langchain\\zhuge\\documents\\company_introduction.txt", encoding="utf-8")
documents = loader.load()

text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
documents = text_splitter.split_documents(documents)

embeddings = HuggingFaceEmbeddings(model_name="D:\\my_models\\sentence-transformers\\all-MiniLM-L6-v2")

vector_store = Chroma.from_documents(documents, embeddings, persist_directory="D:\\vector_store\\chroma_db")

vector_load = Chroma(persist_directory="D:\\vector_stores\\chroma_db", embedding_function=embeddings)

query = "陈李公司是干嘛的？"
result = vector_load.similarity_search_with_score(query)
print(result)
